# Author: chenfan_qu@qcf-568
# CUDA_VISIBLE_DEVICES=9 python infer.py --cfg bisai.py --pth /home/challenge/dataset/workspace/epoch_1.pth
# CUDA_VISIBLE_DEVICES=2,3,4,5,6,7,8,9 python infer.py --cfg bisai.py --pth /home/challenge/dataset/workspace/epoch_1.pth
import os
import cv2
import json
import mmcv
import torch
import pickle
import argparse
import numpy as np
from tqdm import tqdm
from mmengine import ConfigDict
from mmengine.config import Config
from mmdet.utils import register_all_modules
from mmdet.apis import init_detector, inference_detector
from ensemble_boxes import weighted_boxes_fusion  # pip install ensemble_boxes

parser = argparse.ArgumentParser(description="Train a segmentor")
parser.add_argument("--cfg", type=str)  # 推理配置文件, 如bisai.py
parser.add_argument("--pth", type=str)  # 训好的模型pth
parser.add_argument("--sz", type=str, default="800,1333")  # 基本尺度, 这个不用动
args = parser.parse_args()

config_file = args.cfg
checkpoint_file = args.pth


# 获取图片路径，递归遍历文件夹
def get_all_image_paths(root_dir):
    img_paths = []
    for root, _, files in os.walk(root_dir):
        for file in files:
            if file.endswith((".jpg", ".jpeg", ".png")):  # 根据你的图片格式调整
                img_paths.append(os.path.join(root, file))
    return img_paths


def unnorm_box(box, w, h):
    if len(box) == 0:
        return box
    else:
        box[:, 0] = box[:, 0] * w
        box[:, 2] = box[:, 2] * w
        box[:, 1] = box[:, 1] * h
        box[:, 3] = box[:, 3] * h
        return box


def norm_box(box, w, h):
    if len(box) == 0:
        return box
    else:
        # 处理一维数组
        box[0] = box[0] / w  # x1
        box[1] = box[1] / h  # y1
        box[2] = box[2] / w  # x2
        box[3] = box[3] / h  # y2
        return box


register_all_modules()
# build the model from a config file and a checkpoint file
config_file = Config.fromfile(config_file)
config_file.model.test_cfg["rpn"]["nms_pre"] = 5000
config_file.model.test_cfg["rpn"]["max_per_img"] = 5000
config_file.model.test_cfg["rcnn"]["score_thres"] = 0.01
config_file.model.test_cfg["rcnn"]["max_per_img"] = 10
config_file.model.pretrained = args.pth
config_file.model = ConfigDict(**config_file.tta_model, module=config_file.model)

test_data_cfg = config_file.test_dataloader.dataset
while "dataset" in test_data_cfg:
    test_data_cfg = test_data_cfg["dataset"]
if "batch_shapes_cfg" in test_data_cfg:
    test_data_cfg.batch_shapes_cfg = None
test_data_cfg.pipeline = config_file.tta_pipeline
s1, s2 = args.sz.split(",")
assert test_data_cfg.pipeline[1]["transforms"][0][0]["type"] == "Resize"
test_data_cfg.pipeline[1]["transforms"][0][0]["scale"] = (int(s1), int(s2))

model = init_detector(config_file, checkpoint_file, device="cuda:0", cfg_options={})
test_img_dir = "/home/challenge/dataset/testingdata/test_set_A_rename"
for i, file_path in enumerate(tqdm(get_all_image_paths(test_img_dir))):
    img = cv2.imread(file_path)
    h, w = img.shape[:2]

    result = inference_detector(model, img)
    pred_instances = result.pred_instances
    boxes = pred_instances.bboxes.cpu().numpy()
    scores = pred_instances.scores.cpu().numpy()
    labels = pred_instances.labels.cpu().numpy()

    # 正确处理 norm_box
    boxes = [norm_box(box, w, h).tolist() for box in boxes]  # 处理一维数组

    # 加权框融合
    boxes, scores, labels = weighted_boxes_fusion(
        [boxes], 
        [scores], 
        [labels], 
        weights=[1], 
        iou_thr=0.25, 
        skip_box_thr=0.0001,
    )
